{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import os, sys\n",
    "sys.path.append('..')\n",
    "import torch\n",
    "import numpy as np\n",
    "from matplotlib.patches import Circle\n",
    "from skimage.transform import resize\n",
    "import copy\n",
    "import json\n",
    "import matplotlib.pyplot as plt\n",
    "import datetime\n",
    "%matplotlib inline\n",
    "from IPython.display import clear_output\n",
    "\n",
    "from model.net import XNMNet\n",
    "from utils.generate_programs import generate_single_program, load_program_generator\n",
    "from utils.misc import convert_david_program_to_mine, invert_dict, todevice\n",
    "from DataLoader import ClevrDataLoader\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modify variables in this block to your path values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt = '' # checkpoint\n",
    "input_dir = '' # folder that contains your preprocessed question and feature files\n",
    "val_question_pt = os.path.join(input_dir, 'val_questions.pt') # val question file\n",
    "val_feature_pt = os.path.join(input_dir, 'val_features.pt') # val feature file\n",
    "vocab_json = os.path.join(input_dir, 'vocab.json') # vocab file\n",
    "annotation_json = '' # CLEVR_val_scenes.json\n",
    "image_dir = '' # folder that contains CLEVR val images\n",
    "pretrained_dir = '../pretrained/' # should contain david's program generator and vocab, don't change it"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load data and pretrained models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "orig_annotations = json.load(open(annotation_json))['scenes']\n",
    "orig_annotations = { int(s['image_index']):s for s in orig_annotations }\n",
    "\n",
    "loaded = torch.load(ckpt, map_location={'cuda:0': 'cpu'})\n",
    "model_kwargs = loaded['model_kwargs']\n",
    "val_loader_kwargs = {\n",
    "    'question_pt': val_question_pt,\n",
    "    'feature_pt': val_feature_pt,\n",
    "    'vocab_json': vocab_json,\n",
    "    'batch_size': 1,\n",
    "    'shuffle': False\n",
    "}\n",
    "val_loader = ClevrDataLoader(**val_loader_kwargs)\n",
    "vocab = val_loader.vocab\n",
    "\n",
    "model_kwargs.update({'vocab': val_loader.vocab})\n",
    "model = XNMNet(**model_kwargs).to(device)\n",
    "model.load_state_dict(loaded['state_dict'])\n",
    "model.eval()\n",
    "\n",
    "program_generator = load_program_generator(os.path.join(pretrained_dir, 'program_generator.pt')).to(device)\n",
    "david_vocab = json.load(open(os.path.join(pretrained_dir, 'david_vocab.json')))\n",
    "david_vocab['program_idx_to_token'] = invert_dict(david_vocab['program_token_to_idx'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define some functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def question_to_str(question, vocab):\n",
    "    question_str = ' '.join(list(filter(\n",
    "        lambda x: x not in {'<NULL>','<START>','<END>'}, \n",
    "        [vocab['question_idx_to_token'][q.item()] for q in question.squeeze()]\n",
    "    )))\n",
    "    return question_str\n",
    "\n",
    "def predict_david_program(questions, data):\n",
    "    programs, program_inputs = [], []\n",
    "    # generate program using david model for each question\n",
    "    for i in range(questions.size(0)):\n",
    "        question_str = []\n",
    "        for j in range(questions.size(1)):\n",
    "            word = data.vocab['question_idx_to_token'][questions[i,j].item()]\n",
    "            if word == '<START>': continue\n",
    "            if word == '<END>': break\n",
    "            question_str.append(word)\n",
    "        question_str = ' '.join(question_str) # question string\n",
    "        david_program = generate_single_program(question_str, program_generator, david_vocab, device)\n",
    "        david_program = [david_vocab['program_idx_to_token'][i.item()] for i in david_program.squeeze()]\n",
    "        # convert david program to ours. return two index lists\n",
    "        program, program_input = convert_david_program_to_mine(david_program, data.vocab)\n",
    "        programs.append(program)\n",
    "        program_inputs.append(program_input)\n",
    "    # padding\n",
    "    max_len = max(len(p) for p in programs)\n",
    "    for i in range(len(programs)):\n",
    "        while len(programs[i]) < max_len:\n",
    "            programs[i].append(vocab['program_token_to_idx']['<NULL>'])\n",
    "            program_inputs[i].append(vocab['question_token_to_idx']['<NULL>'])\n",
    "    # to tensor\n",
    "    programs = torch.LongTensor(programs).to(device)\n",
    "    program_inputs = torch.LongTensor(program_inputs).to(device)\n",
    "    return programs, program_inputs\n",
    "\n",
    "    \n",
    "\n",
    "def draw_at_coords(ax, coords, attn, img, title, radius=10, target_size=360):\n",
    "    coords = copy.deepcopy(coords)\n",
    "    for i in range(len(coords)):\n",
    "        coords[i] = coords[i][:2]\n",
    "        # [0,28) -> [0,1], (y,x) -> (x,y)\n",
    "        coords[i][0] = coords[i][0] / 27 * target_size\n",
    "        coords[i][1] = coords[i][1] / 27 * target_size\n",
    "        coords[i] = coords[i][::-1]\n",
    "    attn = attn[:len(coords)]\n",
    "    ax.imshow(resize(img, (target_size, target_size)))\n",
    "    ax.axis('off')\n",
    "    ax.set_title(title, fontsize=20)\n",
    "    for coord, a in zip(coords, attn):\n",
    "        a = min(max(a, 0), 1)\n",
    "        patch = Circle(coord, radius, color=(1,1-a,1-a))\n",
    "        ax.add_patch(patch)\n",
    "    \n",
    "\n",
    "def display_helper(ax, attn, title='', target_size=360):\n",
    "    ax.imshow(resize(attn, (target_size, target_size)))\n",
    "    ax.axis('off')\n",
    "    ax.set_title(title, fontsize=20)\n",
    "    \n",
    "\n",
    "def input_score(prompt):\n",
    "    while True:\n",
    "        try:\n",
    "            score = int(input(prompt))\n",
    "        except (KeyboardInterrupt, SystemExit):\n",
    "            raise\n",
    "        except:\n",
    "            print(\"Input must be integer\")\n",
    "            continue\n",
    "        if score in {1,2,3,4}:\n",
    "            break\n",
    "        else:\n",
    "            print(\"Input must be one of (1,2,3,4)\")\n",
    "    return score\n",
    "        \n",
    "\n",
    "def T(x, y, flag=True): # transpose or not\n",
    "    return y, x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# human evaluation of visualization and write average score into the record file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "scores = []\n",
    "total = 100 # how many instances will be visualized\n",
    "cnt = 0\n",
    "sz=6 # size of displayed images\n",
    "\n",
    "for batch in val_loader:\n",
    "    clear_output()\n",
    "    \n",
    "    answers, questions, gt_programs, gt_program_inputs, *batch_input = [todevice(x, device) for x in batch]\n",
    "    answer_str = vocab['answer_idx_to_token'][answers[0].item()]\n",
    "    image_idx = val_loader.idx_cache[0]\n",
    "    annotation = orig_annotations[image_idx]\n",
    "    objects_coord = val_loader.coord_cache[0]\n",
    "    image_file = os.path.join(image_dir, annotation['image_filename'])\n",
    "    img = plt.imread(image_file)\n",
    "    programs, program_inputs = predict_david_program(questions, val_loader)\n",
    "    predict_str, intermediates = model.forward_and_return_intermediates(programs, program_inputs, *batch_input)\n",
    "    if not intermediates:\n",
    "        continue\n",
    "\n",
    "    cnt += 1\n",
    "    # =============== split into branches ================\n",
    "    branches = []\n",
    "    branch = []\n",
    "    intermediates.append(None)\n",
    "    for i in range(len(intermediates)):\n",
    "        if intermediates[i] is None:\n",
    "            if branch:\n",
    "                branches.append(branch)\n",
    "            branch = []\n",
    "        else:\n",
    "            branch.append(intermediates[i])\n",
    "    # ---------------------------------\n",
    "    if len(branches) == 3:\n",
    "        branches[1], branches[2] = branches[2], branches[1]\n",
    "        leg_len = max(len(branches[0]), len(branches[2]))\n",
    "        l = leg_len+len(branches[1])\n",
    "        fig, axs = plt.subplots(*T(3, l), squeeze=False, figsize=T(sz*l, sz*3))\n",
    "        for i in range(3):\n",
    "            for j in range(len(branches[i])):\n",
    "                title, attn = branches[i][j]\n",
    "                c = leg_len+j if i==1 else leg_len-len(branches[i])+j\n",
    "                draw_at_coords(axs[T(i, c)], objects_coord, attn, img, title)\n",
    "            for j in range(l):\n",
    "                if (i!=1 and leg_len-len(branches[i])<=j<leg_len) or\\\n",
    "                    (i==1 and leg_len<=j<leg_len+len(branches[i])):\n",
    "                    continue\n",
    "                axs[T(i,j)].remove()\n",
    "    # ---------------------------------\n",
    "    else:\n",
    "        l = max(len(b) for b in branches)\n",
    "        fig, axs = plt.subplots(*T(len(branches), l), squeeze=False, figsize=T(sz*l, sz*len(branches)))\n",
    "        for i in range(len(branches)):\n",
    "            for j in range(len(branches[i])):\n",
    "                title, attn = branches[i][j]\n",
    "                draw_at_coords(axs[T(i, l-len(branches[i])+j)], objects_coord, attn, img, title)\n",
    "            for j in range(l-len(branches[i])):\n",
    "                axs[T(i,j)].remove()\n",
    "    # ======================================================\n",
    "    \n",
    "\n",
    "    plt.pause(0.0001) # real-time show\n",
    "    print(\"=\"*50)\n",
    "    _, ax = plt.subplots(figsize=(5,5))\n",
    "    display_helper(ax, img, 'Input Image')\n",
    "    plt.pause(0.0001)\n",
    "    print(\" Image Id: %d \\n Question: %s ?\\n Predicted: %s; Answer: %s\" % \n",
    "          (image_idx, question_to_str(questions, vocab), predict_str, answer_str))\n",
    "    \n",
    "    score = input_score(\"%d/%d. Please select one score from: 1(unclear), 2(somewhat unclear), 3(mostly clear) and 4(clear)\"\\\n",
    "                  % (cnt, total))\n",
    "    scores.append(score)\n",
    "    if cnt >= total:\n",
    "        break\n",
    "\n",
    "clear_output()\n",
    "print(\"mean score of %d samples: %.3f\" % (total, np.mean(scores)))\n",
    "worker = input(\"Thank you! Please input your name:\")\n",
    "filename = './data/human_eval_results.txt'\n",
    "with open(filename, 'a') as f:\n",
    "    content = '[Task1] Time:[%s] Worker: %s. Samples: %d. Mean score: %.3f\\n' %\\\n",
    "        (datetime.datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\"), worker, total, np.mean(scores))\n",
    "    f.write(content)\n",
    "    print(\"Write into %s\" % filename)\n",
    "print(\"All Finshed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## show original image of target index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = 0\n",
    "from IPython.display import Image, display\n",
    "\n",
    "for batch in val_loader:\n",
    "    answers, questions, gt_programs, gt_program_inputs, *batch_input = [todevice(x, device) for x in batch]\n",
    "    answer_str = vocab['answer_idx_to_token'][answers[0].item()]\n",
    "    image_idx = val_loader.idx_cache[0]\n",
    "    if image_idx != target:\n",
    "        continue\n",
    "    annotation = orig_annotations[image_idx]\n",
    "    objects_coord = [obj['pixel_coords'] for obj in annotation['objects']]\n",
    "    image_file = os.path.join(image_dir, annotation['image_filename'])\n",
    "    break\n",
    "Image(filename=image_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
